package com.zf.websocket.interceptor;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import javax.servlet.http.HttpSession;
import java.util.Map;

/**
 * WebSocket握手拦截器，自带的有
 * - OriginHandshakeInterceptor：校验是否允许跨域；默认启用
 * - HttpSessionHandshakeInterceptor：处理 HttpSession；默认不使用
 * <p>
 * 自定义 拦截器，在配置中注册后才生效
 */
public class MyHandShakeInterceptor implements HandshakeInterceptor {
    /**
     * 握手前调用的方法
     *
     * @param request    当前请求
     * @param response   响应
     * @param wsHandler  对应访问地址的 WebSocket 处理器
     * @param attributes 与 WebSocketSession 关联的属性，每个 WebSocketSession 都有一个单独的属性
     * @return 是（true）继续握手  还是（false）中止
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {

        // 获取 HttpSession
        HttpSession session = getSession(request);
        if (session != null) {
            // 将 HttpSession 存入 WebSocketSession 的属性中
            attributes.put(HttpSession.class.getName(), session);
        }
        return true;
    }

    private HttpSession getSession(ServerHttpRequest request) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
            return serverRequest.getServletRequest().getSession();
        }
        return null;
    }

    /**
     * 握手完成后调用
     *
     * @param request   当前请求
     * @param response  当前响应
     * @param wsHandler 对应访问地址的 WebSocket 处理器
     * @param exception 握手期间是否出现异常，如果没有就是 null
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {

    }
}
